import pymysql
import time
import uuid

if __name__ == '__main__':

    # 以下参数需要修改成当前训练环境对应的内容，以便进行训练和写入数据库
    cp = "aistudio"
    gpu = "Tesla V100-SXM2-32GB"
    my_model = '/home/aistudio/work/yolov10-main/ultralytics/cfg/models/v10/yolov10n.yaml'
    my_data = '/home/aistudio/work/yolov10-main/data.yaml'
    imgsz = 640
    epochs = 5
    batch = 64
    workers = 2
    device = '0'
    optimizer = 'Adam'
    close_mosaic = 3
    resume = False
    uuid_str = str(uuid.uuid4())
    project = 'runs/train'
    name = 'train_20250212_' + uuid_str
    single_cls = False
    cache = True
    sava_dir = project + '/' + name
    start_time = time.time()

    # model.load('yolov8n.pt') # 加载预训练权重,改进或者做对比实验时候不建议打开，因为用预训练模型整体精度没有很明显的提升
    model = YOLOv10(model=r'/home/aistudio/work/yolov10-main/ultralytics/cfg/models/v10/yolov10n.yaml')
    model.train(data=r'/home/aistudio/work/yolov10-main/data.yaml',
                imgsz=imgsz,
                epochs=epochs,
                batch=batch,
                workers=workers,
                device=device,
                optimizer=optimizer,
                close_mosaic=close_mosaic,
                resume=resume,
                project='runs/train',
                name=name,
                single_cls=single_cls,
                cache=cache,
                )

    for i in range(1000):
        print(i)

    end_time = time.time()
    # 计算训练结束时间
    time_delta = end_time - start_time

    hours = int(time_delta // 3600)
    minutes = int((time_delta % 3600) // 60)
    seconds = int(time_delta % 60)
    # 格式化为符合 TIME 类型的字符串
    train_time = f"{hours:02d}:{minutes:02d}:{seconds:02d}"

    # 建立数据库连接
    conn = pymysql.connect(
        host='chengshaohao.mysql.rds.aliyuncs.com',
        user='root',
        password='Csh@9560',
        database='yolov10',
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )

    try:
        with conn.cursor() as cursor:
            # 定义插入数据的 SQL 语句
            sql = "INSERT INTO train (cp, gpu, model, data, imgsz, epochs, batch, workers, device, optimizer, close_mosaic, resume, single_cls, cache, sava_dir, train_time) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)"
            # 要插入的数据
            val = (cp, gpu, my_model, my_data, imgsz, epochs, batch, workers, device, optimizer, close_mosaic, resume,
                   single_cls, cache, sava_dir, train_time)
            # 执行插入语句
            cursor.execute(sql, val)

        # 提交事务，使插入操作生效
        conn.commit()
        print(cursor.rowcount, "写入完成")
    except Exception as e:
        print(f"Error: {e}")
    finally:
        # 关闭数据库连接
        conn.close()
